import os

from zhipuai import ZhipuAI

client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])

def add(a, b):
    # 加法函数
    return a + b

def subtract(a, b):
    # 减法函数
    return a - b

def multiply(a, b):
    # 乘法函数
    return a * b

def divide(a, b):
    # 除法函数
    return a / b

#定义用于ChatGLM4用于function calling的工具函数声明
tools = [
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "加法函数",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {
                        "type": "number",
                        "description": "第一个数"
                    },
                    "b": {
                        "type": "number",
                        "description": "第二个数"
                    }
                },
                "required": ["a", "b"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "subtract",
            "description": "减法函数",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {
                        "type": "number",
                        "description": "第一个数"
                    },
                    "b": {
                        "type": "number",
                        "description": "第二个数"
                    }
                },
                "required": ["a", "b"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "multiply",
            "description": "乘法函数",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {
                        "type": "number",
                        "description": "第一个数"
                    },
                    "b": {
                        "type": "number",
                        "description": "第二个数"
                    }
                },
                "required": ["a", "b"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "divide",
            "description": "除法函数",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {
                        "type": "number",
                        "description": "第一个数"
                    },
                    "b": {
                        "type": "number",
                        "description": "第二个数"
                    }
                },
                "required": ["a", "b"]
            }
        }
    }
]

prompt = "请帮我计算一下1250908070 乘以 42525463 等于几"
messages = [
    {"role": "user", "content": prompt}
]
response = client.chat.completions.create(
    model="glm-4",  # 填写需要调用的模型名称
    messages= messages,
    tools= tools,
    tool_choice= "auto"
)

if len(response.choices[0].message.tool_calls) >0:

    # 获取函数名称
    function_name = response.choices[0].message.tool_calls[0].function.name
    # 获取函数参数
    function_args = response.choices[0].message.tool_calls[0].function.arguments
    # 将函数参数转换为字典
    function_args_dict = eval(function_args)
    # 调用函数
    function_result = eval(function_name)(**function_args_dict)
    # 将函数结果转换为字符串
    function_result_str = str(function_result)
    #messages.append({"role": "assistant","content": function_result_str})
    messages.append({
        "role": "tool",
        "name": response.choices[0].message.tool_calls[0].function.name,
        "content": prompt + "?结果:" + function_result_str + "\n" + prompt
    })

    print(messages)

    response = client.chat.completions.create(
        model="glm-4",  # 填写需要调用的模型名称
        messages=messages,
    )

    if response.choices[0].message.content:
        print(response.choices[0].message.content)
    else:
        print("No response from ChatGLM4.")
else:
    print(response.choices[0].message.content)